import streamlit as st
from agents.generator import GeneratorAgent
from agents.evaluator import EvaluatorAgent
from agents.memory import MemoryModule
import json

# 加载基准数据
def load_benchmark():
    try:
        with open('data/benchmark.json', 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception:
        return []

benchmark = load_benchmark()

# Streamlit 侧边栏：大模型配置
with st.sidebar:
    st.header("大模型配置")
    llm_api_key = st.text_input("API Key", type="password", value="sk-aLPyvKcaNQehEdTpAOBS467ixdpPdtFzUkA8b7HbOhH1Gp2L")
    llm_model = st.selectbox("模型名称", ["moonshot-v1-8k", "moonshot-v1-32k", "gpt-3.5-turbo", "gpt-4", "自定义模型"])
    llm_temperature = st.slider("采样温度", 0.0, 1.0, 0.7, 0.05)
    llm_base_url = st.text_input("API Base URL", value="https://api.moonshot.cn/v1")

# 初始化智能体，传递大模型配置
generator = GeneratorAgent(api_key=llm_api_key, model=llm_model, temperature=llm_temperature, base_url=llm_base_url)
evaluator = EvaluatorAgent(benchmark)
memory = MemoryModule()

st.title("检索式生成-验证对抗博弈平台")

user_input = st.text_area("请输入检索需求（如：量子材料相关文献）")

# 初始化session_state
if 'round' not in st.session_state:
    st.session_state['round'] = 0
if 'history' not in st.session_state:
    st.session_state['history'] = []
if 'last_feedback' not in st.session_state:
    st.session_state['last_feedback'] = None
if 'finished' not in st.session_state:
    st.session_state['finished'] = False

max_rounds = 5

if st.button("生成下一轮检索式", disabled=st.session_state['finished'] or not user_input):
    if st.session_state['round'] < max_rounds:
        query = user_input
        last_feedback = st.session_state['last_feedback']
        retrieval = generator.generate(query, last_feedback)
        score, feedback = evaluator.evaluate(retrieval)
        # 专家修正输入
        expert_feedback = st.text_input(f"专家修正意见（可选，第{st.session_state['round']+1}轮）", key=f"expert_{st.session_state['round']}")
        if expert_feedback:
            feedback += "\n[专家修正] " + expert_feedback
        # 记录历史
        st.session_state['history'].append({
            'round': st.session_state['round']+1,
            'retrieval': retrieval,
            'score': score,
            'feedback': feedback
        })
        st.session_state['last_feedback'] = feedback
        st.session_state['round'] += 1
        if score >= 0.9 or expert_feedback or st.session_state['round'] >= max_rounds:
            st.session_state['finished'] = True

# 展示历史轮次
for item in st.session_state['history']:
    st.markdown(f"### 第{item['round']}轮")
    st.write("生成检索式：", item['retrieval'])
    st.write("评估分数：", item['score'])
    st.write("评估意见：", item['feedback'])

if st.session_state['finished']:
    st.success("博弈终止，检索式已达标或达到最大轮次！")

if st.button("重置博弈"):
    st.session_state['round'] = 0
    st.session_state['history'] = []
    st.session_state['last_feedback'] = None
    st.session_state['finished'] = False